from peewee import *
import math
import random

db = SqliteDatabase('rocket.db')

ORBIT_DV = 5000 # 常量
GRAVITY = 9.8 # 常量
DOWNSTAGE_TWR = 1.1 # 常量，标记第一级需要的最小twr
UPSTAGE_TWR = 0.1 # 常量，入轨最后一级需要的最小twr
HIGHEST_AIR = 3000 # 常量，到达大气外所需的DV
VAC_FACTOR = 1.2 # 常量，真空环境发动机较地面有多少的性能提升

# 所有的质量单位为千克，推力单位为吨
# 注意:所有的类想要初始化都请使用post_init()!!! 若使用这个方法，请记得加多一点null=True

class BaseModel(Model): # 将所有的Meta类给提前定义好了
    class Meta:
        database = db

class KindOfFuel(BaseModel):
    name = CharField(unique=True)
    density = FloatField()
    price_each_volume = FloatField()
    base_lsp = FloatField()

    def count_volume(self, tons):
        return tons / self.density
    def count_price(self, tons):
        return self.price_each_volume * self.count_volume(tons)
    
class KindOfTank(BaseModel):
    name = CharField(unique=True)
    mass_each_volume = FloatField()
    price_each_volume = FloatField()

    def count_mass(self, volume):
        return volume * self.mass_each_volume
    def count_price(self, volume):
        return self.price_each_volume * volume
    
    class Meta:
        database = db
    
class KindOfEngineCycle(BaseModel):
    name = CharField(unique=True)
    lsp_factor_max = FloatField() # 比例系数
    lsp_factor_min = FloatField()
    price_factor_max = FloatField() # 同上
    price_factor_min = FloatField()
    develop_skill = FloatField() # 初始为0, 之后随着研发次数增加而增长
    develop_time = FloatField() # 你研发所需要的时间
    mass_factor = FloatField()

    def count_everything(self, kofuel, tech_price, force):
        # 返回值：一个四重字典[性能，生产时间，价格，引擎质量]
        # 算法用脚写的
        # tech_price为一个三元字典[性能投入，成本投入, 速度投入],总共为1
        # 都是基础量，可能需要在后面加一个修正
        lsp_factor = self.lsp_factor_min + (self.lsp_factor_max - self.lsp_factor_min) * tech_price['tech']
        price_factor = self.price_factor_max - (self.lsp_factor_max - self.lsp_factor_min) * tech_price['price']
        lsp = kofuel.base_lsp * lsp_factor
        price = pow(force, 1.2) * price_factor / min(math.log(self.develop_skill + math.e), 2)
        time = self.develop_time * (2 - tech_price['time']) / 2
        mass = pow(force, 2) * self.mass_factor
        return {'lsp': lsp, 'price': price, 'time': time, 'mass': mass}
    
    def develop_it(self):
        self.develop_skill += 1

class Engine(BaseModel):
    name = CharField(unique=True, default='default')
    mass = FloatField(null=True) # 单位为kg
    lsp = FloatField(null=True) # 单位为秒
    force = FloatField(null=True) # 单位为ton
    base_price = FloatField(null=True) # 单位为星辰币（？）
    skilled = FloatField(null=True) # 熟练度
    develop_time = FloatField(null=True) # 生产一台发动机所需要的时间
    kind_of_fuel = ForeignKeyField(model=KindOfFuel, backref='engines')
    kind_of_cycle = ForeignKeyField(model=KindOfEngineCycle, backref='engines')

    def post_init(self, name=None, tech_price=None, force=None, kofuel=None, kocycle=None):
        # tech_price为一个三元字典[性能投入，成本投入, 速度投入],总共为1
        # super().__init__(*args, **kwargs) # 
        # super().__init__(self, *args, **kwarg) # 豆包说的
        self.kind_of_fuel = KindOfFuel.select().where(KindOfFuel.id == kofuel)
        self.kind_of_cycle = KindOfEngineCycle.select().where(KindOfEngineCycle.id == kocycle)
        self.force = force
        lsp_price_time = self.kind_of_cycle.count_everything(kofuel=self.kind_of_fuel, tech_price=tech_price, force=force)
        self.mass = lsp_price_time['mass']
        self.lsp = lsp_price_time['lsp']
        self.develop_time = lsp_price_time['time']
        self.skilled = min(0.95 + self.kind_of_cycle.develop_skill * 0.005, 0.99)
        self.name = name
        self.base_price = lsp_price_time['price']

    def after_launch(self, result):
        if result: # 成功
            new_skilled = 1 - (1 - self.skilled) * 0.8
            self.skilled = new_skilled
        else: # 失败
            new_skilled = 1 - (1 - self.skilled) * 0.5
            self.skilled = new_skilled

    def is_success(self):
        result = random.random()
        if result >= self.skilled:
            return False
        return True
    
    def ask_for_price(self):
        return self.base_price * 0.95 / self.skilled

# 发动机组，用于火箭分级的过渡式模型
class EngineGroup(BaseModel):
    # lsp = FloatField()
    total_lsp = FloatField(default=0)
    total_force = FloatField(default=0)
    total_price = FloatField(default=0)
    total_mass = FloatField(default=0)

    def count_everything(self):
        self.total_price = 0
        self.total_force = 0
        self.total_lsp = 0
        self.total_mass = 0
        for engine_t in self.engines:
            self.total_force += engine_t.engine.force * engine_t.amount
            self.total_lsp += engine_t.engine.lsp * engine_t.amount * engine_t.engine.force
            self.total_price += engine_t.engine.ask_for_price() * engine_t.amount
            self.total_mass += engine_t.engine.mass * engine_t.amount
        real_lsp = self.total_lsp / self.total_force
        self.total_lsp = real_lsp

    def launch_mock(self):
        # 这个用来模拟发射时的故障，以后有机会再写
        pass

class EngineTeam(BaseModel):
    engine = ForeignKeyField(model=Engine)
    amount = IntegerField()
    owner = ForeignKeyField(model=EngineGroup, on_delete='CASCADE', backref='engines')

class Tank(BaseModel):
    name = CharField(unique=True)
    kind_of_fuel = ForeignKeyField(model=KindOfFuel)
    kind_of_tank = ForeignKeyField(model=KindOfTank)
    volume = FloatField(null=True)
    # 下面所有数据都需要再加一个修正才能使用（涉及模拟经营元素）
    empty_mass = FloatField(null=True)
    fuel_mass = FloatField(null=True)
    empty_price = FloatField(null=True)
    fuel_price = FloatField(null=True)

    def post_init(self, name=None, kofuel=None, kotank=None, tons=None):
        # super().__init__(self, *args, **kwarg)
        self.name = name
        self.kind_of_fuel = KindOfFuel.select().where(KindOfFuel.id == kofuel)
        self.kind_of_tank = KindOfTank.select().where(KindOfTank.id == kotank)
        self.volume = self.kind_of_fuel.count_volume(tons)
        self.empty_mass = self.kind_of_tank.count_mass(self.volume)
        self.fuel_mass = tons
        self.empty_price = self.kind_of_tank.count_price(self.volume)
        self.fuel_price = self.kind_of_fuel.count_price(tons)

class RocketStage(BaseModel):
    name = CharField(unique=True)
    engine_group = ForeignKeyField(model=EngineGroup)
    tank = ForeignKeyField(model=Tank)
    empty_mass = FloatField(null=True)
    total_mass = FloatField(null=True)
    empty_price = FloatField(null=True)
    total_price = FloatField(null=True)

    def post_init(self, name=None, engine_group=None, tank=None):
        # super().__init__(self, *args, **kwarg)
        self.name = name
        self.engine_group = engine_group
        self.tank = Tank.select().where(Tank.id == tank)
        self.empty_mass = self.engine_group.total_mass + self.tank.empty_mass
        self.total_mass = self.empty_mass + self.tank.fuel_mass
        self.empty_price = self.engine_group.total_price + self.tank.empty_price
        self.total_price = self.empty_price + self.tank.fuel_price

class Rocket(BaseModel):
    name = CharField(unique=True)
    total_mass = FloatField(default=0)
    total_price = FloatField(default=0)
    empty_mass = FloatField(default=0)
    empty_price = FloatField(default=0)

    def get_stage(self):
        return (RocketStage.select().join(StageAndRocket).join(Rocket)
                .where(StageAndRocket.rocket == self).order_by(StageAndRocket.series.asc()))
    
    def count_everything(self):
        self.total_mass = 0
        self.total_price = 0
        self.empty_mass = 0
        self.empty_price = 0
        stages = self.get_stage()
        for stage in stages:
            self.total_mass += stage.total_mass
            self.total_price += stage.total_price
            self.empty_mass += stage.empty_mass
            self.empty_price += stage.empty_price

    def check_dv(self, payload_mass):
        # 给定载荷质量，返回DV
        # 用for遍历stage应当是下面级在前
        mass = payload_mass
        stages = self.get_stage()
        total_dv = 0
        for stage in stages:
            mass += stage.total_mass
        series = 1
        for stage in stages:
            empty_mass = mass - stage.total_mass + stage.empty_mass
            stage_dv = stage.engine_group.total_lsp * 9.8 * math.log(mass / empty_mass)
            if total_dv < ORBIT_DV:
                min_twr = ORBIT_DV / pow(DOWNSTAGE_TWR - UPSTAGE_TWR, 2) * pow(total_dv - ORBIT_DV, 2) + UPSTAGE_TWR
                # 用二次函数模拟
                if min_twr > stage.engine_group.total_force * 1000 / mass:
                    return {'result':1, 'reason':[series, min_twr, stage.engine_group.total_force / mass]}
            if total_dv < HIGHEST_AIR:
                stage_dv *= pow(VAC_FACTOR, total_dv / HIGHEST_AIR)
            total_dv += stage_dv
            series += 1
        return {'result':0, 'reason':total_dv}

class StageAndRocket(BaseModel):
    # 木有感情的中间表
    stage = ForeignKeyField(model=RocketStage, backref='rocket_stage')
    rocket = ForeignKeyField(model=Rocket, backref='prin_rocket')
    series = IntegerField() # 表示火箭是第几级，越大越处于上面级
    class Meta:
        database = db
        indexes = ((('stage', 'rocket'), True),)
